from typing import List, Dict, Any
from loguru import logger
import json
from pathlib import Path


class ScriptGenerator:
    """脚本代码生成器"""
    
    def __init__(self):
        self.indent = "    "
    
    def generate_script(self, actions: List[Dict[str, Any]], script_name: str = "test_script") -> str:
        """从动作列表生成Playwright脚本代码"""
        try:
            script_lines = []
            
            # 添加导入语句
            script_lines.extend(self._generate_imports())
            script_lines.append("")
            
            # 添加测试函数
            script_lines.append(f"async def {script_name}():")
            script_lines.append(f"{self.indent}\"\"\"自动生成的测试脚本\"\"\"")
            script_lines.append("")
            
            # 添加浏览器初始化
            script_lines.extend(self._generate_browser_setup())
            script_lines.append("")
            
            # 生成动作代码
            script_lines.extend(self._generate_actions_code(actions))
            script_lines.append("")
            
            # 添加清理代码
            script_lines.extend(self._generate_cleanup())
            script_lines.append("")
            
            # 添加主函数
            script_lines.extend(self._generate_main_function(script_name))
            
            return "\n".join(script_lines)
            
        except Exception as e:
            logger.error(f"生成脚本失败: {e}")
            raise
    
    def _generate_imports(self) -> List[str]:
        """生成导入语句"""
        return [
            "import asyncio",
            "from playwright.async_api import async_playwright",
            "from loguru import logger",
            "import time",
            "from datetime import datetime",
            "from pathlib import Path"
        ]
    
    def _generate_browser_setup(self) -> List[str]:
        """生成浏览器初始化代码"""
        return [
            f"{self.indent}# 确保截图目录存在",
            f"{self.indent}Path('screenshots').mkdir(exist_ok=True)",
            f"{self.indent}",
            f"{self.indent}async with async_playwright() as p:",
            f"{self.indent}{self.indent}# 启动浏览器",
            f"{self.indent}{self.indent}browser = await p.chromium.launch(headless=False)",
            f"{self.indent}{self.indent}context = await browser.new_context()",
            f"{self.indent}{self.indent}page = await context.new_page()",
            f"{self.indent}{self.indent}",
            f"{self.indent}{self.indent}try:"
        ]
    
    def _generate_actions_code(self, actions: List[Dict[str, Any]]) -> List[str]:
        """生成动作代码"""
        lines = []
        double_indent = self.indent + self.indent + self.indent
        
        for i, action in enumerate(actions):
            action_type = action.get("type")
            data = action.get("data", {})
            
            # 添加注释
            lines.append(f"{double_indent}# 动作 {i+1}: {action_type}")
            
            if action_type == "navigate":
                url = data.get("url", "")
                lines.append(f"{double_indent}await page.goto('{url}')")
                
            elif action_type == "click":
                selector = data.get("selector", "")
                if selector:
                    lines.append(f"{double_indent}await page.click('{selector}')")
                else:
                    x, y = data.get("x", 0), data.get("y", 0)
                    lines.append(f"{double_indent}await page.mouse.click({x}, {y})")
                    
            elif action_type == "input":
                selector = data.get("selector", "")
                value = data.get("value", "")
                if selector:
                    lines.append(f"{double_indent}await page.fill('{selector}', '{value}')")
                    
            elif action_type == "keyboard":
                key = data.get("key", "")
                if key:
                    lines.append(f"{double_indent}await page.keyboard.press('{key}')")
            
            # 添加等待时间
            lines.append(f"{double_indent}await page.wait_for_timeout(500)")
            
            # 为重要动作添加截图
            if action_type in ["click", "navigate"]:
                timestamp = "datetime.now().strftime('%Y%m%d_%H%M%S_%f')[:-3]"
                screenshot_path = f"'screenshots/{action_type}_{{{timestamp}}}.png'"
                lines.append(f"{double_indent}# 截图")
                lines.append(f"{double_indent}try:")
                lines.append(f"{double_indent}{self.indent}await page.screenshot(path={screenshot_path})")
                lines.append(f"{double_indent}except Exception as screenshot_error:")
                lines.append(f"{double_indent}{self.indent}logger.warning(f'截图失败: {{screenshot_error}}')")
            
            lines.append("")
        
        return lines
    
    def _generate_cleanup(self) -> List[str]:
        """生成清理代码"""
        return [
            f"{self.indent}{self.indent}except Exception as e:",
            f"{self.indent}{self.indent}{self.indent}logger.error(f'脚本执行失败: {{e}}')",
            f"{self.indent}{self.indent}{self.indent}raise",
            f"{self.indent}{self.indent}finally:",
            f"{self.indent}{self.indent}{self.indent}await browser.close()"
        ]
    
    def _generate_main_function(self, script_name: str) -> List[str]:
        """生成主函数"""
        return [
            "if __name__ == '__main__':",
            f"{self.indent}asyncio.run({script_name}())"
        ]
    
    def add_assertion(self, script_code: str, assertion_type: str, 
                     selector: str = "", expected_value: str = "", 
                     line_number: int = -1) -> str:
        """在脚本中添加断言"""
        try:
            lines = script_code.split("\n")
            double_indent = self.indent + self.indent + self.indent
            
            # 生成断言代码
            assertion_code = self._generate_assertion_code(
                assertion_type, selector, expected_value, double_indent
            )
            
            # 插入断言
            if line_number > 0 and line_number < len(lines):
                lines.insert(line_number, assertion_code)
            else:
                # 在cleanup之前插入
                for i, line in enumerate(lines):
                    if "except Exception as e:" in line:
                        lines.insert(i - 1, assertion_code)
                        break
            
            return "\n".join(lines)
            
        except Exception as e:
            logger.error(f"添加断言失败: {e}")
            raise
    
    def _generate_assertion_code(self, assertion_type: str, selector: str, 
                               expected_value: str, indent: str) -> str:
        """生成断言代码"""
        if assertion_type == "text_equals":
            return f"{indent}# 断言: 文本等于 '{expected_value}'\n" \
                   f"{indent}actual_text = await page.text_content('{selector}')\n" \
                   f"{indent}assert actual_text == '{expected_value}', f'文本不匹配: 期望 \"{expected_value}\", 实际 \"{{actual_text}}\"'"
        
        elif assertion_type == "text_contains":
            return f"{indent}# 断言: 文本包含 '{expected_value}'\n" \
                   f"{indent}actual_text = await page.text_content('{selector}')\n" \
                   f"{indent}assert '{expected_value}' in actual_text, f'文本不包含: 期望包含 \"{expected_value}\", 实际 \"{{actual_text}}\"'"
        
        elif assertion_type == "element_visible":
            return f"{indent}# 断言: 元素可见\n" \
                   f"{indent}assert await page.is_visible('{selector}'), f'元素不可见: {selector}'"
        
        elif assertion_type == "element_hidden":
            return f"{indent}# 断言: 元素隐藏\n" \
                   f"{indent}assert not await page.is_visible('{selector}'), f'元素应该隐藏: {selector}'"
        
        elif assertion_type == "url_equals":
            return f"{indent}# 断言: URL等于 '{expected_value}'\n" \
                   f"{indent}current_url = page.url\n" \
                   f"{indent}assert current_url == '{expected_value}', f'URL不匹配: 期望 \"{expected_value}\", 实际 \"{{current_url}}\"'"
        
        elif assertion_type == "url_contains":
            return f"{indent}# 断言: URL包含 '{expected_value}'\n" \
                   f"{indent}current_url = page.url\n" \
                   f"{indent}assert '{expected_value}' in current_url, f'URL不包含: 期望包含 \"{expected_value}\", 实际 \"{{current_url}}\"'"
        
        else:
            return f"{indent}# 未知断言类型: {assertion_type}"
    
    def save_generated_script(self, script_code: str, file_path: str) -> None:
        """保存生成的脚本代码"""
        try:
            Path(file_path).parent.mkdir(parents=True, exist_ok=True)
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(script_code)
            
            logger.info(f"生成的脚本已保存到: {file_path}")
            
        except Exception as e:
            logger.error(f"保存生成的脚本失败: {e}")
            raise